// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/tool/trace/filter_project_replayer.h>
#include <pollux/exec/trace_util.h>
#include <pollux/plan/plan_builder.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

namespace kumo::pollux::tool::trace {
core::PlanNodePtr FilterProjectReplayer::createPlanNode(
    const core::PlanNode* node,
    const core::PlanNodeId& nodeId,
    const core::PlanNodePtr& source) const {
  if (node->name() == "Filter") {
    const auto* filterNode = dynamic_cast<const core::FilterNode*>(node);
    POLLUX_CHECK(
        !isFilterProject(dynamic_cast<const core::FilterNode*>(node)),
        "If the target node is a FilterNode, it must be a standalone FilterNode");

    // A standalone FilterNode.
    return std::make_shared<core::FilterNode>(
        nodeId, filterNode->filter(), source);
  }

  const auto* projectNode = dynamic_cast<const core::ProjectNode*>(node);

  // A standalone ProjectNode.
  if (node->sources().empty() || node->sources().front()->name() != "Filter") {
    return std::make_shared<core::ProjectNode>(
        nodeId, projectNode->names(), projectNode->projections(), source);
  }

  // A ProjectNode with a FilterNode as its source.
  // -- ProjectNode [nodeId]
  //   -- FilterNode [nodeId - 1]
  const auto originalFilterNode =
      std::dynamic_pointer_cast<const core::FilterNode>(
          node->sources().front());
  const auto filterNode = std::make_shared<core::FilterNode>(
      nodeId, originalFilterNode->filter(), source);
  const auto projectNodeId = planNodeIdGenerator_->next();
  return std::make_shared<core::ProjectNode>(
      projectNodeId,
      projectNode->names(),
      projectNode->projections(),
      filterNode);
}

bool FilterProjectReplayer::isFilterProject(
    const core::PlanNode* filterNode) const {
  const auto* projectNode =
      dynamic_cast<const core::ProjectNode*>(core::PlanNode::findFirstNode(
          planFragment_.get(), [this](const core::PlanNode* node) {
            return node->id() == std::to_string(std::stoull(nodeId_) + 1);
          }));
  return projectNode != nullptr && projectNode->name() == "Project" &&
      projectNode->sources().size() == 1 &&
      projectNode->sources().front()->id() == nodeId_;
}
} // namespace kumo::pollux::tool::trace
